local Gaussian = {}
Gaussian.__index = Gaussian

function Gaussian:new(mu, sigma, pi, tau)
    pi = pi or 0
    tau = tau or 0
    if mu then
        if not sigma then
            error('sigma argument is needed')
        elseif sigma == 0 then
            error('sigma^2 should be greater than 0')
        end
        pi = sigma ^ -2
        tau = pi * mu
    end
    local self = setmetatable({}, Gaussian)
    self.pi = pi
    self.tau = tau
    return self
end

function Gaussian:mu()
    return self.pi and self.tau / self.pi or 0
end

function Gaussian:sigma()
    return self.pi and math.sqrt(1 / self.pi) or math.huge
end

function Gaussian:__mul(other)
    local pi, tau = self.pi + other.pi, self.tau + other.tau
    return Gaussian(nil, nil, pi, tau)
end

function Gaussian:__div(other)
    local pi, tau = self.pi - other.pi, self.tau - other.tau
    return Gaussian(nil, nil, pi, tau)
end

function Gaussian:__eq(other)
    return self.pi == other.pi and self.tau == other.tau
end

function Gaussian:__lt(other)
    return self:mu() < other:mu()
end

function Gaussian:__le(other)
    return self:mu() <= other:mu()
end

function Gaussian:__gt(other)
    return self:mu() > other:mu()
end

function Gaussian:__ge(other)
    return self:mu() >= other:mu()
end

function Gaussian:__tostring()
    local args = {self:mu(), self:sigma()}
    return ('N(mu=%.3f, sigma=%.3f)'):format(table.unpack(args))
end

setmetatable(Gaussian, {__call = Gaussian.new})

local Matrix = {}
Matrix.__index = Matrix

function Matrix:new(src)
    local self = setmetatable({}, Matrix)
    if type(src) == 'table' then
        local num = nil
        local rectangular = true
        local everything_are_number = true
        for _, row in ipairs(src) do
            if num == nil then
                num = #row
            elseif #row ~= num then
                rectangular = false
                goto exit_loop
            end
            for _, item in ipairs(row) do
                if type(item) ~= 'number' then
                    everything_are_number = false
                    goto exit_loop
                end
            end
        end
        ::exit_loop::
        if not rectangular or not everything_are_number then
            error('src must be a rectangular array of numbers')
        end
        self.array = src
    else 
        error('src must be a table')
    end 
    return self 
end 

function Matrix:height()
    return #self.array 
end 

function Matrix:width()
    return #self.array[1]
end 

function Matrix:transpose()
    local height, width = self:height(), self:width()
    local src = {}
    for c = 1, width do
        src[c] = {}
        for r = 1, height do
            src[c][r] = self.array[r][c]
        end
    end
    return Matrix(src)
end 

function Matrix:minor(row_n, col_n)
    local height, width = self:height(), self:width()
    if row_n < 0 or row_n >= height then 
        error(('row_n should be between 0 and %d'):format(height))
    elseif col_n < 0 or col_n >= width then 
        error(('col_n should be between 0 and %d'):format(width))
    end 
    local src = {}
    for r = 0, height - 1 do 
        if r  == row_n then goto continue end 
        local row = {}
        table.insert(src, row)
        for c = 0, width -1 do 
            if c == col_n then goto continue end 
            table.insert(row, self.array[r + 1][c + 1])
            ::continue::
        end 
        ::continue::
    end 
    return Matrix(src)
end 

function Matrix:determinant()
    local height, width = self:height(), self:width()
    if height ~= width then
        error('Only square matrix can calculate a determinant')
    end
    local det = 1
    local swaps = 0
    local tmp = setmetatable({},Matrix)
    for r = 1 , height do
        tmp[r] = {}
        for c = 1, width do
            tmp[r][c] = self.array[r][c]
        end
    end
    for r = 1, height do
        -- 寻找主元
        local max_element = math.abs(tmp[r][r])
        local max_row = r
        for i = r + 1, height do
            if math.abs(tmp[i][r]) > max_element then
                max_element = math.abs(tmp[i][r])
                max_row = i
            end
        end
        -- 如果主元为0，则矩阵奇异，行列式为0
        if max_element == 0 then
            return 0
        end
        -- 交换当前行和最大行
        if r ~= max_row then
            tmp[r], tmp[max_row] = tmp[max_row], tmp[r]
            swaps = swaps + 1
        end
        -- 消元
        for i = r + 1, height do
            local c = -tmp[i][r] / tmp[r][r]
            for j = r, height do
                tmp[i][j] = tmp[i][j] + c * tmp[r][j]
            end
        end
        det = det * tmp[r][r]
    end
    if swaps % 2 == 1 then
        det = -det
    end
    return det
end

function Matrix:adjugate()
    local height, width = self:height(), self:width()
    if height ~= width then 
        error('Only square matrix can be adjugated')
    end 
    if height == 2 then 
        local a, b = self.array[1][1], self.array[1][2]
        local c, d = self.array[2][1], self.array[2][2]
        return Matrix({{d, -b}, {-c, a}})
    end 
    local src = {}
    for r = 1, height do 
        src[r] = {}
        for c = 1, width do 
            local sign = ((r + c - 2) % 2 == 0) and 1 or -1
            src[r][c] =
                self:minor(r - 1, c - 1):determinant() * sign
        end 
    end 
    return Matrix(src):transpose()
end 

function Matrix:inverse()
    if self:height() == 1 and self:width() == 1 then 
        return Matrix({{1. / self.array[1][1]}})
    end 
    return self:adjugate() * (1. / self:determinant()) 
end 

function Matrix:__add(other)
    local height, width = self:height(), self:width()
    if height ~= other:height() or width ~= other:width() then 
        error('Must be same size')
    end 
    local src = {}
    for r = 1, height do 
        src[r] = {}
        for c = 1, width do 
            src[r][c] =
                self.array[r][c] + other.array[r][c]
        end 
    end 
    return Matrix(src, height, width)
end 

function Matrix:__mul(other)
    if type(other) == 'number' then 
        local height, width = self:height(), self:width()
        local src = {}
        for r = 1, height do 
            src[r] = {}
            for c = 1, width do 
                src[r][c] =
                    other * self.array[r][c]
            end 
        end 
        return Matrix(src)
    elseif getmetatable(other) == Matrix then 
        if self:width() ~= other:height() then 
            error('Bad size')
        end 
        local height, width =
            self:height(), other:width()
        local src = {}
        for r = 1, height do 
            src[r] = {}
            for c = 1, width do 
                local sum = 0
                for x in ipairs(self.array[1]) do 
                    sum =
                        sum +
                        (self.array[r][x] *
                         other.array[x][c])
                end 
                src[r][c] =
                    sum
            end 
        end 
        return Matrix(src)
    else 
        error('The operand should be a number or a matrix')
    end 
end 

function Matrix:__tostring()
    local args = {type(self), tostring(table.concat(self.array, ', '))}
    return ('%s(%s)'):format(table.unpack(args))
end

setmetatable(Matrix, {__call = Matrix.new})

return {
  Gaussian=Gaussian,
  Matrix=Matrix,
}